{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 使用 Prim 算法生成最小生成树\n",
    "\n",
    "Prim 算法使用的是一种类似 \"两军对垒\" 的形式，使用 MST 性质，生成最小生成树。\n",
    "\n",
    "算法步骤：\n",
    "1. 选取任意节点作为初始节点，将顶点集划分为 S 和 V-S 两个子集。\n",
    "2. 选取 $i\\in S, j\\in V-S$, 当 $c[i,j]$ 是连接两子集的最小边，那么将 $V[j]$ 加入到集合 $S$ 中。\n",
    "3. 当 $S = V$ 时，算法结束，否则继续步骤 2。\n",
    "\n",
    "算法时间复杂度为: $O(VlgV + ElgV) = O(ElgV)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 初始化\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 邻接矩阵用 NumPy 存储\n",
    "def AdjacencyMatrix(Path):\n",
    "    '''\n",
    "    从 Path 中读取邻接矩阵.\n",
    "    Input:\n",
    "    - Path: 邻接矩阵保存地址.\n",
    "    \n",
    "    Return:\n",
    "    - E: 邻接矩阵.\n",
    "    - V: 顶点集合.\n",
    "    '''\n",
    "    E = None\n",
    "    V = None\n",
    "    index = 0\n",
    "    with open(Path,'r') as f:\n",
    "        while True:\n",
    "            data = f.readline()\n",
    "            if len(data) == 0:\n",
    "                break\n",
    "            # 删除换行符\n",
    "            data = data.rstrip('\\n')\n",
    "            # 分词\n",
    "            data = data.split('\\t')\n",
    "            if index == 0:\n",
    "                E = np.zeros((len(data),len(data)))\n",
    "                E[index] = np.array(data,dtype=int)\n",
    "                V = np.array(list(range(len(data))))\n",
    "            else:\n",
    "                E[index] = np.array(data,dtype=int)\n",
    "            index = index + 1\n",
    "    return E,V\n",
    "    \n",
    "Path = 'example.txt'\n",
    "# 读取邻接矩阵和顶点集\n",
    "E,V = AdjacencyMatrix(Path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def NonArgMin(E,V):\n",
    "    '''\n",
    "    找非零最小值,且不在顶点集内.\n",
    "    Input:\n",
    "    - E: 邻接矩阵.\n",
    "    - V: 标记过的顶点集.\n",
    "    \n",
    "    Return:\n",
    "    - ind: 最小值下标.\n",
    "    '''\n",
    "    ind = None\n",
    "    mini = 1e15\n",
    "    for i in V:\n",
    "        for j in range(E.shape[0]):\n",
    "            if E[i,j] < mini and E[i,j]!=0 and (j not in V):\n",
    "                mini = E[i,j]\n",
    "                ind = [i,j]\n",
    "    return ind"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def Prim(E):\n",
    "    '''\n",
    "    使用 Prim 算法生成最小生成树。\n",
    "    Input:\n",
    "    - E: 邻接矩阵.\n",
    "    \n",
    "    Return:\n",
    "    - nE: 新的邻接矩阵，代表最小生成树的连接关系，只有 0 和 1.\n",
    "    '''\n",
    "    nE = np.zeros_like(E,dtype=int)\n",
    "    Length = [] # 生成树大小\n",
    "    V = [0] # 选取初始顶点\n",
    "    for i in range(E.shape[0]-1):\n",
    "        #print(i)\n",
    "        ind = NonArgMin(E,V)\n",
    "        x = ind[0]\n",
    "        y = ind[1]\n",
    "        V.append(y)\n",
    "        Length.append(E[x,y])\n",
    "        nE[x,y] = nE[y,x] = 1\n",
    "    return nE,sum(Length)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1888.0\n",
      "[0]\n",
      "[0 0]\n",
      "[0 0 0]\n",
      "[1 0 0 0]\n",
      "[0 0 0 0 0]\n",
      "[0 1 0 0 1 0]\n",
      "[0 0 0 0 0 1 0]\n",
      "[0 1 1 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 1 0]\n",
      "[0 0 0 0 0 1 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 1 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 1 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 1 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 1 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 1 1 1 0 0 0 0]\n",
      "[0 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 1 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 1 0 0 0 0 0 1 0 0 0 0 0 0 0 0 0 0]\n"
     ]
    }
   ],
   "source": [
    "nE,length = Prim(E)\n",
    "print(length)\n",
    "# 以符合题目要求的形式输出\n",
    "for i in range(E.shape[0]):\n",
    "    print(nE[i,:i+1])\n",
    "    \n",
    "# 输出到文件中\n",
    "with open('prim_output.txt','w') as f:\n",
    "    for i in range(E.shape[0]):\n",
    "        for j in range(i+1):\n",
    "            if j!=i:\n",
    "                f.write('%s\\t' % nE[i,j])\n",
    "            else:\n",
    "                f.write('%s\\n' % nE[i,j])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch",
   "language": "python",
   "name": "torch"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
